use std::collections::HashMap;

use solutiont2::{json::JsonParser, ufset::UnionFindSet};

static FILE_PATH: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/district.json");

pub fn count_provinces() -> String {
    let object = JsonParser::from_file(FILE_PATH).unwrap().parse().unwrap();
    let n = object.map.len();
    object
        .map
        .into_iter()
        .fold(vec!["0".to_string(); n], |mut result, (k, v)| {
            let k: usize = k.parse().unwrap();
            let map = &v[0].expect_object().map;
            let mut indices = map
                .keys()
                .enumerate()
                .map(|(idx, key)| (key.as_str(), idx))
                .collect::<HashMap<_, _>>();
            let mut ufset = UnionFindSet::new(map.len());
            for (i, (key, value)) in map.iter().enumerate() {
                for v in value.iter().flat_map(|v| &v.expect_array().0) {
                    let next_index = indices.len();
                    let &mut j = indices.entry(v.expect_string()).or_insert(next_index);
                    ufset.union(i, j);
                }
            }
            for (x, i) in indices {
                println!("{x:?}: {i}/{}", ufset.find(i));
            }
            result[k - 1] = ufset.classes().to_string();
            result
        })
        .join(",")
}
